from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager,
)
import pandas as pd
from datetime import date, time, datetime, timedelta
from vnpy.trader.constant import Status, Interval, Direction


cpi_path = 'E:\\GitCloneProject\\quant25\\01Src\\vnpy\\environment_evaluation\\data\\cpi_data.csv'


class CpiRsiStrategy(CtaTemplate):
    """"""
    # 策略参数
    author = "Youpeng Hu"
    ma_length_short = 5   # 短期MA周期
    ma_length_long = 20  # 长期MA周期
    rsi_length = 14     # RSI周期
    cpi_threshold = 0.5  # CPI阈值
    fixed_size = 1       # 固定交易手数
    
    # 数据存储
    cpi_data = {}
    last_cpi_month = None
    current_month = None
    
    parameters = [
        "ma_length_short",
        "ma_length_long",
        "rsi_length",
        "cpi_threshold",
        "fixed_size"
    ]
    
    variables = [
        "ma_short",
        "ma_long",
        "rsi_value",
        "cpi_condition_met",
        "position_signal"
    ]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        self.bg = BarGenerator(
            self.on_bar,
            interval=Interval.DAILY,
            daily_end=time(15, 0)
        )
        self.am = ArrayManager(size=10)
        self.load_cpi_data()  # 加载CPI数据

    def load_cpi_data(self):
        """加载本地CPI数据"""
        df = pd.read_csv(cpi_path)
        self.cpi_data = df.set_index("month").to_dict()["nt_yoy"]
        self.write_log(f"成功加载CPI数据，覆盖月份：{list(df['month'])}")

    def on_init(self):
        """初始化策略"""
        self.load_bar(10)  # 初始化历史数据
        self.write_log("策略初始化完成")

    def on_bar(self, bar: BarData):
        """交易逻辑"""
        # 1. 更新指标数据
        self.am.update_bar(bar)
        if not self.am.inited:
            return
        
        # 计算MA指标
        self.ma_short = self.am.sma(self.ma_length_short, array=True)[-1]
        self.ma_long = self.am.sma(self.ma_length_long, array=True)[-1]
        
        # 计算RSI指标
        self.rsi_value = self.am.rsi(self.rsi_length)
        
        # 2. 检查CPI条件
        current_date = bar.datetime.strftime("%Y%m")
        self.cpi_condition_met = False
        print(type(self.cpi_data))
        if current_date in self.cpi_data:
            print("in")
            last_month = (datetime.strptime(current_date, "%Y%m") - pd.DateOffset(months=1)).strftime("%Y%m")
            if last_month in self.cpi_data:
                self.cpi_condition_met = self.cpi_data[last_month] < self.cpi_threshold
        else:
            print("not in")
        # 3. 生成交易信号
        ma_cross = self.ma_short > self.ma_long
        signal = ""
        
        if ma_cross and self.cpi_condition_met:
            signal = "BUY"
        elif not ma_cross or not self.cpi_condition_met:
            signal = "SELL"
        
        # 4. 执行交易
        if signal and self.pos == 0:
            price = bar.close_price
            if signal == "BUY":
                self.buy(price + 1e-4, self.fixed_size)
            else:
                self.sell(price - 1e-4, self.fixed_size)

    def on_order(self, order: OrderData):
        """订单更新处理"""
        pass

    def on_trade(self, trade: TradeData):
        """成交处理"""
        pass

    def on_stop(self):
        """策略停止"""
        self.write_log("策略停止")
    